from torch import nn
from typing import Tuple


class ResidualBlock(nn.Module):
    def __init__(self, in_channels: int, out_channels: int, strides: Tuple[int, ...] = (1, )):
        super(ResidualBlock, self).__init__()
        self._conv1 = nn.Conv2d(in_channels, out_channels, kernel_size=(3, 3), padding=1, stride=strides)
        self._conv2 = nn.Conv2d(out_channels, out_channels, kernel_size=(3, 3), padding=1)
        if in_channels != out_channels or strides != (1, ):
            self._conv3 = nn.Conv2d(in_channels, out_channels, kernel_size=(1,), stride=strides)
        else:
            self._conv3 = None
        self._bn1 = nn.BatchNorm2d(num_features=out_channels)
        self._bn2 = nn.BatchNorm2d(num_features=out_channels)
        self._act = nn.ReLU(inplace=True)

    def forward(self, x):
        y = self._act(self._bn1(self._conv1(x)))
        y = self._bn2(self._conv2(y))
        return self._act(y + (x if not self._conv3 else self._conv3(x)))


class ResNet(nn.Module):
    def __init__(self):
        super(ResNet, self).__init__()
        self._stage_1 = nn.Sequential(
            nn.Conv2d(1, 64, kernel_size=(7, ), stride=(2, ), padding=3),
            nn.BatchNorm2d(64), nn.ReLU(),
            nn.MaxPool2d(kernel_size=3, stride=2, padding=1)
        )
        self._stage_2 = nn.Sequential(
            ResidualBlock(64, 64),
            ResidualBlock(64, 64)
        )
        self._stage_3 = nn.Sequential(
            ResidualBlock(64, 128, strides=(2, )),
            ResidualBlock(128, 128)
        )
        self._stage_4 = nn.Sequential(
            ResidualBlock(128, 256, strides=(2,)),
            ResidualBlock(256, 256)
        )
        self._stage_5 = nn.Sequential(
            ResidualBlock(256, 512, strides=(2,)),
            ResidualBlock(512, 512)
        )
        self._avg_pool = nn.AdaptiveAvgPool2d((1, 1))
        self._flatten = nn.Flatten()
        self._linear = nn.Linear(512, 10)

    def forward(self, x):
        x = self._stage_1(x)
        x = self._stage_2(x)
        x = self._stage_3(x)
        x = self._stage_4(x)
        x = self._stage_5(x)
        x = self._avg_pool(x)
        x = self._flatten(x)
        x = self._linear(x)
        return x
